//SPDX-FileCopyrightText: Copyright 2022-2024 深圳市同心圆网络有限公司
//SPDX-License-Identifier: GPL-3.0-only

package impl

import (
	"os"

	"atomgit.com/openlinksaas-org/proxy/ssh_pkg/gen_code"
	"golang.org/x/crypto/ssh"
)

type ChainClient struct {
	realClient *ssh.Client
	parent     *ChainClient
}

func NewChainClient(parent *ChainClient, authInfo *gen_code.AuthInfo) (*ChainClient, error) {
	authList := []ssh.AuthMethod{}

	if authInfo.PrivKeyPath != "" {
		keyData, err := os.ReadFile(authInfo.PrivKeyPath)
		if err != nil {
			return nil, err
		}
		var singer ssh.Signer
		if authInfo.PrivKeyPassword == "" {
			singer, err = ssh.ParsePrivateKey(keyData)
			if err != nil {
				return nil, err
			}
		} else {
			singer, err = ssh.ParsePrivateKeyWithPassphrase(keyData, []byte(authInfo.PrivKeyPassword))
			if err != nil {
				return nil, err
			}
		}
		authList = append(authList, ssh.PublicKeys(singer))
	}
	if authInfo.Password != "" {
		authList = append(authList, ssh.Password(authInfo.Password))
	}

	cfg := &ssh.ClientConfig{
		User:            authInfo.Username,
		Auth:            authList,
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
	}

	if parent == nil {
		c, err := ssh.Dial("tcp", authInfo.Addr, cfg)
		if err != nil {
			return nil, err
		}
		return &ChainClient{
			realClient: c,
			parent:     nil,
		}, nil
	} else {
		conn, err := parent.realClient.Dial("tcp", authInfo.Addr)
		if err != nil {
			return nil, err
		}
		sc, nc, rc, err := ssh.NewClientConn(conn, authInfo.Addr, cfg)
		if err != nil {
			conn.Close()
			return nil, err
		}
		c := ssh.NewClient(sc, nc, rc)
		return &ChainClient{
			realClient: c,
			parent:     parent,
		}, nil
	}
}

func (c *ChainClient) Close() {
	c.realClient.Close()
	c.realClient = nil
	if c.parent != nil {
		c.parent.Close()
		c.parent = nil
	}
}
